import uuid

from fastapi import APIRouter, HTTPException, Depends
from pydantic import BaseModel

from auth_service.auths import verify_token
from tables.knowledge import KG, KnowledgeModel
from tables.user import Users
from tables.vector import Vectors
from utils.main import str_to_hash
from datetime import datetime
router = APIRouter(
    prefix="/rag/knowledge",
    tags=["knowledge management"],
    responses={404: {"description": "Not found"}}
)


class CreateKnowledge(BaseModel):
    rag_name: str
    description: str = None

@router.post("/create", summary="创建知识库")
async def create_knowledge_api(
        body: CreateKnowledge,
        current_user = Depends(verify_token)
):
    id = str(uuid.uuid4())
    try:
        now = datetime.now()
        # 格式化为 yyyy-MM-dd HH:mm:ss
        formatted_date_time = now.strftime('%Y-%m-%d %H:%M:%S')
        res = KG.create_knowledge(current_user.id, id, body.rag_name, body.description,formatted_date_time)
        if not res:
            raise Exception("知识库创建失败")
        Vectors.create_collection(collection_name=str_to_hash(id), is_base=True)
        return {"code": 200, "message": "知识库创建成功"}
    except Exception as e:
        # 删除创建的知识库
        KG.delete_knowledge(current_user.id, id)
        raise HTTPException(status_code=500, detail=f"知识库创建失败: {str(e)}")

@router.delete("/delete", summary="删除知识库")
async def delete_knowledge_api(
        id: str,
        current_user = Depends(verify_token)
):
    try:
        kg = KG.get_knowledge(current_user.id, id)
        res = KG.delete_knowledge(current_user.id, id)
        if not res:
            raise Exception("知识库删除失败")
        if kg.is_set:
            Vectors.delete_collection(str_to_hash(id))
        return {"code": 200, "message": "知识库删除成功"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"知识库删除失败: {str(e)}")


class UpdateKnowledge(BaseModel):
    id: str
    embed: str
    reranker: str
    llm: str
    rag_name: str = ""
    description: str = ""
    chunk_size: int = 0
    chunk_overlap: int = 0
    chunk_type: int = 1
    rag_type: int = 1
    content_augmentation: int = 2
    retrieval_type: int = 3
    mix_threshold: float = 0.75
    top_k: int = 5
    is_set: bool = True
    update_date:str = ""

@router.post("/update", summary="更新知识库")
async def update_knowledge_api(
        body: UpdateKnowledge,
        current_user = Depends(verify_token)
):
    try:
        old_kg = KG.get_knowledge(current_user.id, body.id)
        if not body.is_set:
            pass
        elif any([body.chunk_type != old_kg.chunk_type, body.embed != old_kg.embed, body.rag_type != old_kg.rag_type]):
            Vectors.delete_collection(str_to_hash(body.id))
            Vectors.create_collection(collection_name=str_to_hash(body.id), is_base=True)
        body.is_set = True
        now = datetime.now()
        # 格式化为 yyyy-MM-dd HH:mm:ss
        formatted_date_time = now.strftime('%Y-%m-%d %H:%M:%S')
        body.update_date = formatted_date_time

        res = KG.update_knowledge(current_user.id, body.id, old_kg, body)
        if not res:
            raise Exception("知识库更新失败")
        return {"code": 200, "message": "知识库更新成功"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"知识库更新失败: {str(e)}")


@router.get("/get_knowledge", summary="获取知识库")
async def get_knowledge_api(
        rag_id: str,
        current_user = Depends(verify_token)
):
    try:
        res = KG.get_knowledge(userid=current_user.id, id=rag_id)
        return {"code": 200, "message": "知识库获取成功", "data": res}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"知识库获取失败: {str(e)}")


@router.get("/get_all_knowledge", summary="获取知识库列表")
async def get_all_knowledge_api(
        query: str = None,
        current_user = Depends(verify_token)
):
    try:
        res = KG.get_all_knowledge(query,userid=current_user.id)
        return {"code": 200, "message": "知识库列表获取成功", "data": res}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"知识库列表获取失败: {str(e)}")


@router.get("/get_public_knowledge", summary="获取公共知识库列表")
async def get_public_knowledge_api(
        current_user = Depends(verify_token)
):
    try:
        public_user = Users.get_user_by_name("public")
        res = KG.get_public_knowledge(userid=public_user.id)
        return {"code": 200, "message": "知识库列表获取成功", "data": res}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"知识库列表获取失败: {str(e)}")